import argparse, torch

# 获取设备信息
if torch.cuda.is_available():
    device = torch.device("cuda:0")
    use_gpu = True
else:
    device = torch.device("cpu")
    use_gpu = False

# ----参数设置----
dataset = 'python'
max_epoch = 300
batch_size = 32
strategy = "v" # 'v', 'n', 'm(默认)'
# 模型种类需要在"main.py"中设置
# Transformer层
num_layers = 4
n_head = 4
d_model = 360  # d_model建议不超过50

# ----打印参数----
print("数据集:", dataset)
print("epoch:", max_epoch)
print("batch:", batch_size)
print("---融合层---")
print("策略:",strategy)
print("---Transformer---")
print("num_layers:", num_layers)
print("n_head:", n_head)
print("d_model:", d_model)

parser = argparse.ArgumentParser()


def str2bool(v):
    return v.lower() in ('true')


def add_argument_group(name):
    arg = parser.add_argument_group(name)
    return arg


net_arg = add_argument_group('Network')
# 1）嵌入层
net_arg.add_argument('--fix_gaz_emb', type=str2bool, default=True)
net_arg.add_argument('--gaz_dropout', type=float, default=0.5)
# 2-1）"BLSTM_GAT_CRF"独有
net_arg.add_argument('--lstm_layer', type=int, default=1)
net_arg.add_argument('--bilstm_flag', type=str2bool, default=True)
net_arg.add_argument('--dropout', type=float, default=0.1)
net_arg.add_argument('--droplstm', type=float, default=0.5)
# 2-2）"Transformer_GAT_CRF"独有
#   1.Transformer层参数
parser.add_argument('--num_layers', default=num_layers)
parser.add_argument('--n_head', default=n_head)
parser.add_argument('--d_model', default=d_model)  # d_model建议不超过50
# 3）GAT
net_arg.add_argument('--gat_nhidden', type=int, default=180) # 单个注意机制的规模
net_arg.add_argument('--gat_nhead', type=int, default=4)
net_arg.add_argument('--gat_layer', type=int, default=1, choices=[1, 2])
net_arg.add_argument("--alpha", type=float, default=0.1)
net_arg.add_argument('--dropgat', type=float, default=0.5)
# 4）融合层
net_arg.add_argument('--strategy', type=str, default=strategy, choices=['v', 'n', 'm'])

# 数据集
import pathes

if dataset == 'python':
    train_file_path = pathes.python_ner_path + "/PY.train"
    test_file_path = pathes.python_ner_path + "/PY.test"
    dev_file_path = pathes.python_ner_path + "/PY.dev"
    tagscheme = "BIO"
elif dataset == 'resume':
    train_file_path = pathes.resume_ner_path + "/train.char.bmes"
    test_file_path = pathes.resume_ner_path + "/test.char.bmes"
    dev_file_path = pathes.resume_ner_path + "/dev.char.bmes"
    tagscheme = "BMES"
# Data
data_arg = add_argument_group('Data')
data_arg.add_argument('--dataset_name', type=str, default=dataset)
data_arg.add_argument('--tagscheme', type=str, default=tagscheme)
data_arg.add_argument('--train_file', default=train_file_path)
data_arg.add_argument('--test_file', default=test_file_path)
data_arg.add_argument('--dev_file', default=dev_file_path)
data_arg.add_argument('--gaz_file', default=pathes.yangjie_rich_pretrain_word_path)
data_arg.add_argument('--char_embedding_path', default=pathes.yangjie_rich_pretrain_unigram_path)
data_arg.add_argument('--data_stored_directory', type=str, default="./data/generated_data/")
data_arg.add_argument('--param_stored_directory', type=str, default="./data/model_param/")

preprocess_arg = add_argument_group('Preprocess')
preprocess_arg.add_argument('--norm_char_emb', type=str2bool, default=False)
preprocess_arg.add_argument('--norm_gaz_emb', type=str2bool, default=True)
preprocess_arg.add_argument('--number_normalized', type=str2bool, default=True)
preprocess_arg.add_argument('--max_sentence_length', type=int, default=250)

learn_arg = add_argument_group('Learning')
learn_arg.add_argument('--batch_size', type=int, default=batch_size)
learn_arg.add_argument('--max_epoch', type=int, default=max_epoch)
learn_arg.add_argument('--lr', type=float, default=0.001)
learn_arg.add_argument('--lr_decay', type=float, default=0.01)
learn_arg.add_argument('--use_clip', type=str2bool, default=False)
learn_arg.add_argument('--clip', type=float, default=5.0)
learn_arg.add_argument("--optimizer", type=str, default="Adam", choices=['Adam', 'SGD'])
learn_arg.add_argument("--l2_penalty", type=float, default=0.00000005)
# Misc
misc_arg = add_argument_group('Misc')
misc_arg.add_argument('--refresh', type=str2bool, default=False)
misc_arg.add_argument('--use_gpu', type=str2bool, default=use_gpu)
misc_arg.add_argument('--visible_gpu', type=int, default=0)  # 指定能够使用的GPU
misc_arg.add_argument('--random_seed', type=int, default=1)


def get_args():
    args, unparsed = parser.parse_known_args()
    if len(unparsed) > 1:
        print("Unparsed args: {}".format(unparsed))
    return args, unparsed
